package com.cfp4cloud.cfp.knowledge.support.handler.rag.strategy.impl;

import cn.hutool.core.util.BooleanUtil;
import cn.hutool.core.util.NumberUtil;
import com.cfp4cloud.cfp.knowledge.dto.AiSliceReCallDTO;
import com.cfp4cloud.cfp.knowledge.dto.AiSliceReCallRequestDTO;
import com.cfp4cloud.cfp.knowledge.entity.AiDatasetEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiSliceEntity;
import com.cfp4cloud.cfp.knowledge.support.handler.rag.strategy.RagHelper;
import com.cfp4cloud.cfp.knowledge.support.handler.rag.strategy.UnifiedRagStrategy;
import com.cfp4cloud.cfp.knowledge.support.provider.ModelProvider;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.DimensionAwareEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

/**
 * 标准检索策略实现
 * <p>
 * 适用于大多数向量数据库的标准检索逻辑 支持向量检索、全文检索和重排序功能
 * </p>
 *
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class StandardRetrievalStrategy implements UnifiedRagStrategy {

	private final ModelProvider modelProvider;

	@Override
	public boolean supports(String handlerType) {
		// 标准检索策略支持所有处理类型
		return true;
	}

	/**
	 * 执行知识检索操作
	 * @param requestDTO 检索请求参数
	 * @param dataset 数据集实体
	 * @param embeddingStore 向量存储
	 * @return 检索结果列表
	 */
	@Override
	public List<AiSliceReCallDTO> performRetrieval(AiSliceReCallRequestDTO requestDTO, AiDatasetEntity dataset,
			EmbeddingStore<TextSegment> embeddingStore) {

		log.debug("使用标准检索策略进行知识召回: {}", requestDTO.getContent());

		// 1. 向量化查询内容
		DimensionAwareEmbeddingModel embeddingModel = modelProvider.getEmbeddingModel(dataset.getEmbeddingModel());
		Embedding queryEmbedding = embeddingModel.embed(requestDTO.getContent()).content();

		// 2. 执行向量相似度搜索
		List<EmbeddingMatch<TextSegment>> vectorMatches = performVectorSearch(queryEmbedding, requestDTO, dataset,
				embeddingStore);

		// 3. 执行全文检索（如果启用）
		List<EmbeddingMatch<TextSegment>> fullTextMatches = RagHelper.performFullTextSearch(embeddingStore, requestDTO);

		// 4. 合并检索结果
		List<EmbeddingMatch<TextSegment>> mergedMatches = RagHelper.mergeSearchResults(vectorMatches, fullTextMatches);

		// 5. 执行重排序（如果启用）
		List<EmbeddingMatch<TextSegment>> finalMatches = new ArrayList<>(mergedMatches);

		if (BooleanUtil.toBoolean(requestDTO.getReRanking())) {
			finalMatches = RagHelper.performReranking(mergedMatches, requestDTO.getContent());
		}

		// 6. 转换为DTO格式返回
		List<AiSliceReCallDTO> results = finalMatches.stream().map(embeddingMatch -> {
			AiSliceReCallDTO dto = new AiSliceReCallDTO();
			Metadata metadata = embeddingMatch.embedded().metadata();

			dto.setSliceId(metadata.getLong(AiSliceEntity.Fields.id));
			dto.setContent(embeddingMatch.embedded().text());
			dto.setScore(NumberUtil.formatPercent(embeddingMatch.score(), 4));
			dto.setMetadata(metadata.toString());

			return dto;
		}).toList();

		log.debug("标准检索策略: 返回 {} 个结果", results.size());
		return results;
	}

	/**
	 * 执行向量相似度搜索
	 */
	private List<EmbeddingMatch<TextSegment>> performVectorSearch(Embedding queryEmbedding,
			AiSliceReCallRequestDTO requestDTO, AiDatasetEntity dataset, EmbeddingStore<TextSegment> embeddingStore) {
		// 构建搜索请求
		EmbeddingSearchRequest searchRequest = RagHelper.buildSearchRequest(queryEmbedding, requestDTO, dataset);

		// 执行搜索
		EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);

		return searchResult.matches();
	}

}